
from Configurable import *
from Util import *
from urlparse import urlparse
import grp
import os.path
import pwd
import shutil
import subprocess
import sys
import tarfile
import time

class Setup (Configurable):
    
    def __init__(self, conf):
        Configurable.__init__(self, conf)
        self.exit_bad_vm_version = 2
        self.exit_cannot_get_vdi = 4
        self.exit_missing_vdi = 5
        self.exit_no_module = 6
            
    def check_is_root(self):
        user = pwd.getpwuid(os.getuid())[0]
        if user != "root":
            log("The setup must be executed as root. Exiting.")
            sys.exit(os.EX_NOPERM)    
        
    def check_VirtualBox_kernel_module_presence(self):
        cmd = "lsmod | grep vboxdrv"
        mod_exists = shell_exec(cmd)
        if mod_exists == "":
            log("vboxdrv is not present in the current kernel, trying to insert it...")
            mod_insert = "modprobe vboxdrv"
            mod_insert_output = shell_exec(mod_insert)
            if mod_insert_output != "":
                log("Cannot insert VirtualBox kernel module.", 2)
                log("Please check your installation, aborting.", 2)
                sys.exit(os.EX_NOTFOUND)
            else:
                log("vboxdrv successfully inserted, continuing...")
        else:
            log("vboxdrv is present, continuing...")
   
    # Getting which user will be able to launch testcloud:
    def choose_user(self):
        found = False
        while not found:
            user = raw_input("Please enter the login of the user who will run Testcloud: ")
            users = pwd.getpwall()
            found = False
            for u in users:
                if u[0] == user: 
                    found = True
                    return user
                
            if not found:
                log("User %s does not exist!" % user,1)


    def choose_vdis(self):
        print ""
        log("The following virtual disk images are available on the Testcloud server:")
        f = urllib.urlopen(self.vdis_url)
        vdis = eval(f.read())
        counter = 1
        for vdi in vdis:
            print "%d- %s (%s)" % (counter, vdi[0], vdi[1])  
        print
        chosen = False
        while not chosen:
            index = raw_input("Please choose the virtual disk you wish to install for Testcloud: ")
            try:
                ivdi = int(index)
            except ValueError, ve:
                print "Invalid selection."
            for i in range(0, len(vdis)):
                if ivdi == i + 1:
                    chosen = True
                    return vdis[i]
            if not chosen:
                print "Invalid selection."

        return None
    
    def download_vdi(self, vdi):
        vdi_url = vdi[1]
        remote_filename = os.path.split(urlparse(vdi_url).path)[1]
        vdi_download_path = "%s" % os.path.join(self.shared_vdi_dir, remote_filename)
        download_file(vdi_url, vdi_download_path)
        return vdi_download_path
        
        
    def uncompress_vdi(self, vdi_download_path): 
        vdifile = os.path.split(vdi_download_path)[1]
        index = vdifile.find(self.compression_extension)
        vdifile = vdifile[0:index]
        vdifile_path = os.path.join(self.shared_vdi_dir, vdifile)  

        if os.path.exists(vdi_download_path):
            print "Uncompressing %s..." % vdifile
            tar = tarfile.open(vdi_download_path)
            try:
                entry = tar.getmember(vdifile)
                tar.extract(entry, self.shared_vdi_dir)
                os.remove(vdi_download_path)
            except KeyError, e: 
                print "ERROR: %s " % e
                sys.exit(self.exit_cannot_get_vdi)
            tar.close()
            return vdifile_path
            
        else:
            log("No VDI could be found in %s, aborting." % self.shared_vdi_dir,2)
            sys.exit(self.exit_missing_vdi) 

    def chown_chmod_vdi(self, vdifile_path):
        pw = pwd.getpwnam("root")
        ruid = pw.pw_uid
        #testcloud_group = grp.getgrnam("testcloud")
        testcloud_group = grp.getgrnam("root")
        guid = testcloud_group.gr_gid
        os.chown(vdifile_path, ruid, guid)
        os.chmod(vdifile_path, 0444)

    def choose_vdi_memory(self):
        cmd = "free -m | grep 'Mem:' | awk ' { print $2 } '"
        proc = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
        freemem = int(proc.communicate()[0])
        if 255 < freemem < 500:
            log("The station has 512 Mb of available memory or lower. Testcloud will use 256 Mb of memory.")
            use_memory = 256
        elif freemem > 500:
            use_memory = int(freemem / 4)
            log("Testcloud will use half of the available memory: %d Mb" % use_memory)
        else:
            log("Warning: the station's available memory is less than 256 Mb.",1)
            return 0
        
        return use_memory
   
    def install_user_data(self, user, vdifile_path):
        log("Setting up Testcloud for user %s." % user) 
        shell_exec("usermod -G testcloud %s" % user)
        testcloud_home = self.get_user_testcloud_dir(user)
        if not os.path.exists(testcloud_home):
            os.mkdir(testcloud_home)
            os.chown(testcloud_home, pwd.getpwnam(user).pw_uid, grp.getgrnam(user).gr_gid)
        
        if (os.path.exists(vdifile_path)):
            log("Copying %s into %s's Testcloud directory (%s)..." % (os.path.split(vdifile_path)[1], user, testcloud_home))
            user_vdi_path = os.path.join(testcloud_home, os.path.split(vdifile_path)[1])
            shutil.copyfile(vdifile_path, user_vdi_path)       
            #os.chown(user_vdi_path, pwd.getpwnam(user).pw_uid, grp.getgrnam("testcloud").gr_gid)
            os.chown(user_vdi_path, pwd.getpwnam(user).pw_uid, pwd.getpwnam(user).pw_uid)
        else:
            log("Error: the original virtual disk image %s could not be found, aborting." % vdifile_path,2)
            sys.exit(self.exit_missing_vdi)
            
    
    def setup_VirtualBox(self, user, vdifile_path, mem):
        # Warn msg: TODO: how to setup more correctly this virtualbox vm ??
        log("Note that the error below, when it occurs, is just a warning which won't prevent the virtual disk to be configured.") 
        log("[ERROR: Could not find a hard disk with location '$location' in the media registry.]")
        log("")
        
        vdifile_name = os.path.split(vdifile_path)[1]
        vdiname = os.path.splitext(vdifile_name)[0] 
        testcloud_home = self.get_user_testcloud_dir(user)
        log("Creating VM %s... " % vdiname)
        log(shell_exec("su - %s -c \"VBoxManage createvm --name %s --register --basefolder %s\"" % (user, vdiname, testcloud_home)))
        log("Attaching disk %s..." % vdifile_name)
        log(shell_exec("su - %s -c \"VBoxManage modifyvm %s --hda %s/%s\"" % (user, vdiname, testcloud_home, vdifile_name)))
        log(shell_exec("su - %s -c \"VBoxManage modifyvm %s --nic1 nat --nictype1 Am79C973 --cableconnected1 on\"" % (user, vdiname)))
        log("Configuring VM memory (%s)..." % mem)
        log(shell_exec("su - %s -c \"VBoxManage modifyvm %s --memory %dMB\"" % (user, vdiname, mem)))
        log("Turning off audio...")
        log(shell_exec("su - %s -c \"VBoxManage modifyvm %s --audio null\"" % (user, vdiname)))
            

    def setup(self):
        self.check_is_root()
        self.check_VirtualBox_kernel_module_presence()
        vdi_index = self.choose_vdis()
        vdi_download_path = self.download_vdi(vdi_index)
        vdifile_path = self.uncompress_vdi(vdi_download_path)
        self.chown_chmod_vdi(vdifile_path)
        mem = self.choose_vdi_memory()
        user = self.choose_user()
        self.install_user_data(user, vdifile_path)
        self.setup_VirtualBox(user, vdifile_path, mem)
        log("Testcloud has been configured for user %s." % user)
        #print "User %s will have to logout before being allowed to run Testcloud." % user

        


